package com.lee.lure.server.ai.aiclient.milvus;

import io.milvus.client.MilvusClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.*;
import io.milvus.param.dml.SearchParam;

import java.util.List;

/**
 * @author yanwei li create  2024-07-04 14:59
 */
public class MilvusClientWrapper {

    private MilvusConnectionPool milvusConnectionPool;

    public MilvusClientWrapper(MilvusConnectionPool milvusConnectionPool){
        this.milvusConnectionPool = milvusConnectionPool;
    }

    public void createCollections(String collectionName,int dimension){
        MilvusClient client = null;
        try {
            client = milvusConnectionPool.getConnection();
            CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withDescription("Collections Name for business")
                    .withShardsNum(3)
                    .addFieldType(FieldType.newBuilder()
                            .withName("id")
                            .withDescription("Primary Key")
                            .withDataType(DataType.Int64)
                            .withPrimaryKey(true)
                            .withAutoID(true)
                            .build())
                    .addFieldType(FieldType.newBuilder()
                            .withName("embedding")
                            .withDescription("Vector field")
                            .withDataType(DataType.FloatVector)
                            .withDimension(dimension)
                            .build())
                    .build();

            client.createCollection(createCollectionParam);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }finally {
            if (client != null) {
                milvusConnectionPool.releaseConnection(client);
            }
        }

    }

    public SearchResultData search(String collectionName, List<List<Float>> vectors, int topK) {
        MilvusClient client = null;
        try {
            client = milvusConnectionPool.getConnection();
            SearchParam searchParam = SearchParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withMetricType(MetricType.L2)
                    .withTopK(topK)
                    .withVectors(vectors)
                    .withVectorFieldName("embedding")
                    .build();
            final R<SearchResults> search = client.search(searchParam);
            return  search.getData().getResults();

        } catch (InterruptedException e) {
            e.printStackTrace();
            return null;
        } finally {
            if (client != null) {
                milvusConnectionPool.releaseConnection(client);
            }
        }
    }

    public void shutdown() {
        milvusConnectionPool.shutdown();
    }
}
